package org.intrasearch.core.mapreduce;

import com.google.common.collect.*;
import org.slf4j.Logger;

import java.util.Collection;
import java.util.SortedSet;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class MapReduceImpl<I, K extends Comparable<?>, V extends Comparable<?>> implements MapReduce<I, K, V>
{
    private final Object key;
    private final SortedSetMultimap<K, V> mapped = Multimaps.synchronizedSortedSetMultimap(TreeMultimap.<K, V>create());
    private final ExecutorService reduceExecutor = Executors.newSingleThreadExecutor();
    private final ExecutorService mapExecutor = Executors.newSingleThreadExecutor();
    private ReduceTaskEntry reduceTaskEntry;
    private MapTask<I, K, V> mapTask;
    private final Logger logger;
    private final ConcurrentHashMap<K, AtomicInteger> keyToThreshold = new ConcurrentHashMap<K, AtomicInteger>();
    private Integer halfReduceThreshold;

    public MapReduceImpl(Object key, Logger logger)
    {
        this.key = key;
        this.logger = logger;
    }

    public Object getKey()
    {
        return key;
    }

    public void map(final K key, V value)
    {
        mapped.put(key, value);
        checkReduceAndSubmitOnTrigger(key);
        checkHalfReduce();
    }

    private void checkHalfReduce()
    {
        if (halfReduceThreshold == null || reduceTaskEntry == null || mapped.size() < halfReduceThreshold)
        {
            return;
        }
        TreeMultimap<Integer, K> sizeToKey = getValueCollectionSizeToKey();
        int i = 0;
        for (Integer size : sizeToKey.keySet())
        {
            for (K key : sizeToKey.get(size))
            {
                runReduceFor(key);
                i += size;
            }
            if (i >= mapped.size() / 2)
            {
                break;
            }
        }
    }

    private TreeMultimap<Integer, K> getValueCollectionSizeToKey()
    {
        TreeMultimap<Integer, K> valueCollectionSizeToKey = TreeMultimap.create(Ordering.natural().reverse(), Ordering.natural());
        for (K key : ImmutableSet.copyOf(mapped.keySet()))
        {
            valueCollectionSizeToKey.put(mapped.get(key).size(), key);
        }
        return valueCollectionSizeToKey;
    }

    private void checkReduceAndSubmitOnTrigger(final K key)
    {
        if (reduceTaskEntry != null && reduceTaskEntry.shouldTrigger(key, mapped.get(key)))
        {
            runReduceFor(key);
        }
    }

    private void runReduceFor(final K key)
    {
        ReduceRunnable runnable = new ReduceRunnable(key);
        reduceExecutor.submit(runnable);
    }

    public Collection<V> reduce(K key)
    {
        return mapped.removeAll(key);
    }

    public SortedSet<K> keySet()
    {
        return (SortedSet<K>) mapped.keySet();
    }

    public void setReduceTask(ReduceTask<K, V> reduceTask)
    {
        reduceTaskEntry = new ReduceTaskEntry(reduceTask);
    }

    public void forceReduceAll()
    {
        for (K key : ImmutableSet.copyOf(mapped.keySet()))
        {
            runReduceFor(key);
        }
    }

    public void setMapTask(MapTask<I, K, V> mapTask)
    {
        this.mapTask = mapTask;
    }

    public void mapInput(final I input)
    {
        mapExecutor.submit(new Runnable()
        {
            public void run()
            {
                try
                {
                    mapTask.map(input, MapReduceImpl.this);
                }
                catch (Exception e)
                {
                    logger.error("Error while map", e);
                }
            }
        });
    }

    public void finishAndExecute(final Runnable runnable)
    {
        new Thread(new Runnable()
        {
            public void run()
            {
                try
                {
                    shutdownAndWaitTerminate(mapExecutor);
                    forceReduceAll();
                    shutdownAndWaitTerminate(reduceExecutor);
                    runnable.run();
                }
                catch (Exception e)
                {
                    logger.error("Error while finish task", e);
                }
            }
        }).start();
    }

    private void shutdownAndWaitTerminate(ExecutorService executor)
            throws InterruptedException
    {
        executor.shutdown();
        for (;;)
        {
            if (executor.awaitTermination(10, TimeUnit.MINUTES))
            {
                break;
            }
        }
    }

    private void lazyInitThresholdAndLastFuture(K key)
    {
        if (!keyToThreshold.containsKey(key))
        {
            keyToThreshold.putIfAbsent(key, new AtomicInteger(1));
        }
    }

    public void setHalfReduceThreshold(int halfReduceThreshold)
    {
        this.halfReduceThreshold = halfReduceThreshold;
    }

    private class ReduceTaskEntry
    {
        private final ReduceTask<K, V> reduceTask;

        private ReduceTaskEntry(ReduceTask<K, V> reduceTask)
        {
            this.reduceTask = reduceTask;
        }

        public boolean shouldTrigger(K key, Collection<V> collection)
        {
            lazyInitThresholdAndLastFuture(key);
            AtomicInteger threshold = keyToThreshold.get(key);
            int valueCollectionSizeThreshold = threshold.get();
            boolean trigger = collection.size() >= valueCollectionSizeThreshold;
            if (trigger)
            {
                getAndMultiplyBy(threshold, 2);
            }
            return trigger;
        }

        private int getAndMultiplyBy(AtomicInteger atomicInteger, int i)
        {
            for (;;)
            {
                int current = atomicInteger.get();
                if (atomicInteger.compareAndSet(current, current * i))
                {
                    return current;
                }
            }
        }
    }

    private class ReduceRunnable implements Runnable
    {
        private final K key;

        public ReduceRunnable(K key)
        {
            this.key = key;
        }

        public void run()
        {
            final SortedSet<V> reducedCollection = mapped.removeAll(key);
            try
            {
                if (!reducedCollection.isEmpty())
                {
                    reduceTaskEntry.reduceTask.reduce(key, reducedCollection);
                }
            }
            catch (Exception e)
            {
                mapped.putAll(key, reducedCollection);
                logger.error("Error while reduce", e);
            }
        }
    }
}
